
#include "arcfour.h"

/*
code from https://en.wikipedia.org/wiki/RC4
----------------------------------------------------
for i from 0 to 255
  S[i] := i
endfor
j := 0
for i from 0 to 255
  j := (j + S[i] + key[i mod keylength]) mod 256
  swap values of S[i] and S[j]
endfor
 */
export ArcFour *rc4init(const int8 *key, int16 len) {
    ArcFour *p = malloc(sizeof(struct s_arcfour));
    if (p == NULL) {
        fprintf(stderr, "Memory allocation failed: %s\n", strerror(errno));
        exit(EXIT_FAILURE);
    }

    for (int pos = 0; pos < 256; pos++) {
        p->s[pos] = pos;
    }
    p->i = p->j = p->k = 0;
    int8 tmp;
    for (; p->i < 256; p->i++) {
        p->j = (p->j + p->s[p->i] + key[p->i % len]) % 256;
        //swap
        tmp = p->s[p->i];
        p->s[p->i] = p->s[p->j];
        p->s[p->j] = tmp;
    }
    p->i = p->j = 0;
    int32 n;
    rc4whitewash(n, p);
    return p;
}

/*
code from https://en.wikipedia.org/wiki/RC4
----------------------------------------------------
i := (i + 1) mod 256
j := (j + S[i]) mod 256
swap values of S[i] and S[j]
t := (S[i] + S[j]) mod 256
K := S[t]
output K
 */
int8 rc4byte(ArcFour *p) {
    p->i = (p->j + 1) % 256;
    p->j = (p->j + p->s[p->i]) % 256;
    int8 tmp = p->s[p->i];
    p->s[p->i] = p->s[p->j];
    p->s[p->j] = tmp;
    int8 t = (p->s[p->i] + p->s[p->j]) % 256;
    p->k = p->s[t];
    return p->k;
}

export int8 *rc4encrypt(ArcFour *rc4, const int8 *text, int16 len) {
    int8 *cipher = (int8 *) malloc(len + 1);
    if (cipher == NULL) {
        fprintf(stderr, "Memory allocation failed: %s\n", strerror(errno));
        exit(EXIT_FAILURE);
    }

    for (int16 x = 0; x < len; x++) {
        cipher[x] = text[x] ^ rc4byte(rc4);
    }
    return cipher;
}
